import pandas as pd
from MappingTabel import NameToTitle, TitleToNo, SexToNo


def get_title(name):
    str1 = name.split(',')[1]  # Mr. Owen Harris
    str2 = str1.split('.')[0]  # Mr
    str3 = str2.strip()
    return str3


# 填充空缺项
def fill_nan(data):
    data['Age'] = data['Age'].fillna(data.Age.mean())
    data['Fare'] = data['Fare'].fillna(data['Fare'].mean())
    data['Embarked'] = data['Embarked'].fillna('U')
    data['Cabin'] = data['Cabin'].fillna('U')
    return data


def preprocess():
    data = pd.read_csv('./titanic.csv')
    # 填充空缺项
    data = fill_nan(data)

    # 数据处理
    data['Cabin'] = data['Cabin'].map(lambda c: c[0])
    data['Name'] = data['Name'].map(get_title).map(NameToTitle).map(TitleToNo)
    data.drop('Ticket', axis=1, inplace=True)

    # 特征提取
    data['Sex'] = data['Sex'].map(SexToNo)

    PclassDf = pd.get_dummies(data['Pclass'], prefix='Pclass').astype(int)
    data = pd.concat([data, PclassDf], axis=1)
    data.drop('Pclass', axis=1, inplace=True)

    NameDf = pd.get_dummies(data['Name'], prefix='Name').astype(int)
    data = pd.concat([data, NameDf], axis=1)
    data.drop('Name', axis=1, inplace=True)

    CabinDf = pd.get_dummies(data['Cabin'], prefix='Cabin').astype(int)
    data = pd.concat([data, CabinDf], axis=1)
    data.drop('Cabin', axis=1, inplace=True)

    EmbarkedDf = pd.get_dummies(data['Embarked'], prefix='Embarked').astype(int)
    data = pd.concat([data, EmbarkedDf], axis=1)
    data.drop('Embarked', axis=1, inplace=True)

    familyDf = data['SibSp'] + data['Parch'] + 1
    data['FamilyNum'] = familyDf

    pd.set_option('display.width', None)
    pd.set_option('display.max_rows', None)
    # print(data.iloc[:3, :])
    # print(data.shape)
    return data


def get_correlation_matrix(data):
    corr = data.corr()
    print(corr['Survived'].sort_values(ascending=False))


def process_sex(sex):
    return SexToNo[sex]


def process_name(name):
    return TitleToNo[NameToTitle[(get_title(name))]]


def process_cabin(cabin):
    return cabin[0]


if __name__ == '__main__':
    data = preprocess()
    get_correlation_matrix(data)

